source("./function_code/relabel_vars.R")

hazardCoefPlot = function(model1, model2="none", fileID) {

  ## No matter what, get data for first model
  # Coefficients and SEs of first model
  survcoef = summary(model1)$coefficients[,1]
  survse = summary(model1)$coefficients[,4]
  
  # Create data
  survData = data.frame(vars=names(survcoef), coef=survcoef, se=survse)
  survData$lower95 = survData$coef - 1.96*survData$se
  survData$upper95 = survData$coef + 1.96*survData$se
  survData$model = "All"
  survData = survData %>% filter(!grepl("\\:", vars))
  sdata = rbind(survData)
  
  # If there's a second model, get that data too
  if (length(model2) > 1) {
    # Coefficients and SEs of second model
    survcoef2 = summary(model2)$coefficients[,1]
    survse2 = summary(model2)$coefficients[,4]
    
    # Create data
    survData2 = data.frame(vars=names(survcoef2), coef=survcoef2, se=survse2)
    survData2$lower95 = survData2$coef - 1.96*survData2$se
    survData2$upper95 = survData2$coef + 1.96*survData2$se
    survData2$model = "Split"
    survData2 = survData2 %>% filter(!grepl("\\:", vars))
    
    # Put together data from both models
    sdata = rbind(survData, survData2)
  }
  
  # Make variable labels
  sdata = relabelVars(sdata)
  
  # Label the two models
  sdata$model = factor(sdata$model, levels=c("All", "Split"))
  
  # Statistical significance  
  sdata$signif = ifelse(sdata$lower95 > 0 | sdata$upper95 < 0, "Yes", "No")
  
  if (length(model2) > 1) {
    # Generate coefficient plot for two models where model 2 disaggregates negotiations
    hazPlot = ggplot(sdata, aes(coef, fct_rev(vars))) + geom_vline(xintercept=0, linetype=2) + 
      geom_pointrange(aes(xmin=lower95, xmax=upper95, color=model, shape=signif), position=position_dodge(width=-0.6)) + 
      theme_bw() + xlab("Coefficient estimate") +
      scale_y_discrete(name="") + 
      scale_shape_manual("Statistically\nsignificant (95%)", values=c(16, 15)) +
      scale_color_manual("Model", labels=c("All", "Disaggregated"), values=c("gray10", "gray60")) +
      theme(legend.position="bottom", 
            legend.direction = "vertical") +
      guides(color = guide_legend(order=1),
             shape = guide_legend(order=2))
  } else {
    # Generate coefficient plot for just one model that uses all negotiations together
    hazPlot = ggplot(sdata, aes(coef, fct_rev(vars))) + geom_vline(xintercept=0, linetype=2) + 
      geom_pointrange(aes(xmin=lower95, xmax=upper95, shape=signif), position=position_dodge(width=0.6)) + 
      theme_bw() + xlab("Coefficient estimate") +
      scale_y_discrete(name="") + 
      scale_shape_manual("Statistically\nsignificant (95%)", values=c(16, 15)) +
      theme(legend.position="bottom", 
            legend.direction = "vertical") +
      guides(color = guide_legend(order=1),
             shape = guide_legend(order=2))
  }
  
  print(hazPlot)
  
  # Make file name
  fn = paste0("./figures/hazardCoefPlot", fileID, ".pdf")
  
  # Create file
  ggsave(fn, width=4.5, height=7)
}
